import cv2
import numpy as np
from collections import defaultdict, deque
import math


class CollisionDetector:
    def __init__(self, bird_view_width=1980, bird_view_height=1020, debug_enabled=False):
        self.bird_view_width = bird_view_width
        self.bird_view_height = bird_view_height
        self.debug_enabled = debug_enabled
        self.detected_accidents = []  # 存储检测到的事故

        # 碰撞检测参数
        self.min_collision_frames = 3  # 最小连续碰撞帧数
        self.distance_threshold = 50  # 碰撞距离阈值(像素)
        self.speed_threshold = 5  # 相对速度阈值(像素/帧)

        # 车辆状态跟踪
        self.vehicle_history = defaultdict(list)
        self.vehicle_speed = {}
        self.vehicle_direction = {}

        # 碰撞对历史记录（避免重复报告）
        self.collision_pairs = set()

    def update_vehicle_state(self, tracks, frame_width, frame_height):
        """更新车辆状态（位置、速度、方向）"""
        # 计算缩放比例
        scale_x = self.bird_view_width / frame_width
        scale_y = self.bird_view_height / frame_height

        for track_id, track in tracks.items():
            if len(track) < 2:
                continue

            # 转换到鸟瞰图坐标
            current_pos = (track[-1][0] * scale_x, track[-1][1] * scale_y)

            # 存储位置历史
            if track_id not in self.vehicle_history:
                self.vehicle_history[track_id] = deque(maxlen=10)
            self.vehicle_history[track_id].append(current_pos)

            # 计算速度（像素/帧）
            if len(self.vehicle_history[track_id]) >= 2:
                prev_pos = self.vehicle_history[track_id][-2]
                dx = current_pos[0] - prev_pos[0]
                dy = current_pos[1] - prev_pos[1]
                speed = math.sqrt(dx ** 2 + dy ** 2)
                self.vehicle_speed[track_id] = speed

                # 计算方向向量
                if speed > 0.1:  # 避免除以零
                    self.vehicle_direction[track_id] = (dx / speed, dy / speed)

    def detect_accidents(self, frame, tracks, frame_width, frame_height):
        """检测车辆碰撞（追尾、侧碰）"""
        annotated_frame = frame.copy()
        current_accidents = []

        # 更新车辆状态
        self.update_vehicle_state(tracks, frame_width, frame_height)

        # 获取所有车辆ID
        vehicle_ids = list(tracks.keys())

        # 检查所有车辆对
        for i in range(len(vehicle_ids)):
            id1 = vehicle_ids[i]
            if id1 not in self.vehicle_history or len(self.vehicle_history[id1]) < 2:
                continue

            for j in range(i + 1, len(vehicle_ids)):
                id2 = vehicle_ids[j]
                if id2 not in self.vehicle_history or len(self.vehicle_history[id2]) < 2:
                    continue

                # 获取最新位置
                pos1 = self.vehicle_history[id1][-1]
                pos2 = self.vehicle_history[id2][-1]

                # 计算距离
                distance = math.sqrt((pos1[0] - pos2[0]) ** 2 + (pos1[1] - pos2[1]) ** 2)

                # 检查是否可能碰撞
                if distance < self.distance_threshold:
                    # 检查是否已经报告过这对碰撞
                    collision_key = tuple(sorted((id1, id2)))
                    if collision_key in self.collision_pairs:
                        continue

                    # 获取速度向量
                    speed1 = self.vehicle_speed.get(id1, 0)
                    speed2 = self.vehicle_speed.get(id2, 0)
                    dir1 = self.vehicle_direction.get(id1, (0, 0))
                    dir2 = self.vehicle_direction.get(id2, (0, 0))

                    # 计算相对速度
                    rel_speed_x = dir1[0] * speed1 - dir2[0] * speed2
                    rel_speed_y = dir1[1] * speed1 - dir2[1] * speed2
                    rel_speed = math.sqrt(rel_speed_x ** 2 + rel_speed_y ** 2)

                    # 检查相对速度是否超过阈值
                    if rel_speed > self.speed_threshold:
                        # 计算碰撞角度（用于区分碰撞类型）
                        dx = pos2[0] - pos1[0]
                        dy = pos2[1] - pos1[1]
                        angle = math.atan2(dy, dx) * 180 / math.pi

                        # 确定碰撞类型
                        if abs(angle) < 45 or abs(angle) > 135:
                            collision_type = "追尾碰撞"
                        else:
                            collision_type = "侧面碰撞"

                        # 记录碰撞
                        current_accidents.append((f"车辆 {id1} 和 {id2}", collision_type))
                        self.collision_pairs.add(collision_key)

                        # 在原始帧上绘制碰撞标记
                        center_x = int((pos1[0] + pos2[0]) / 2 / (self.bird_view_width / frame_width))
                        center_y = int((pos1[1] + pos2[1]) / 2 / (self.bird_view_height / frame_height))
                        self._draw_collision(annotated_frame, center_x, center_y, collision_type)

        # 存储检测到的事故
        if current_accidents:
            self.detected_accidents.extend(current_accidents)

        return annotated_frame, current_accidents

    def _draw_collision(self, frame, x, y, collision_type):
        """在帧上绘制碰撞标记"""
        cv2.circle(frame, (x, y), 25, (0, 0, 255), 3)
        cv2.putText(frame, collision_type, (x - 50, y - 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
        cv2.putText(frame, "COLLISION!", (x - 40, y),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    def get_accident_history(self):
        return self.detected_accidents

    def reset(self):
        """重置检测器状态"""
        self.detected_accidents = []
        self.vehicle_history.clear()
        self.vehicle_speed.clear()
        self.vehicle_direction.clear()
        self.collision_pairs.clear()